"""
数据层 | 基类
"""

import math

from bson import ObjectId
from pymongo import ASCENDING
from pymongo.collection import Collection
from pymongo.cursor import Cursor
from pymongo.results import (
    InsertManyResult as PyMongoInsertManyResult,
    InsertOneResult as PyMongoInsertOneResult,
    UpdateResult as PyMongoUpdateResult,
)

from db.mongodb_db.client import mongo_db
from db.mongodb_db.transaction import inject_transaction_session
from models import (
    POPage,
    TPO,
    PageQuery,
)
from models.other import (
    InsertManyResult,
    InsertOneResult,
    UpdateResult,
)
from my_typing import *
from utils.data_utils import chunk_iterable
from utils.data_utils.special import ObjectIdStr
from utils.mock_utils import DATETIME

DocumentType = TypeVar("DocumentType", bound=Mapping[str, Any])
"""
MongoDB 的文档类型

References:
    :class:`bson.typings._DocumentType`
"""


class ServiceBase(Generic[TPO]):
    """ 数据层 | 基类 """

    po: type[TPO]
    """ PO 对象 """

    collection_name: str
    """ 集合名称 """

    single_field_index: ListStr = []
    """ 单字段索引 """

    union_field_index: list[ListStr] = []
    """ 联合字段索引 """

    deleted_query: dict[str, bool] = {"is_deleted": True}
    """ 数据库中已删除的查询条件 """

    undeleted_query: dict[str, bool] = {"is_deleted": False}
    """ 数据库中未删除的查询条件 """

    @classmethod
    def get_collection(cls) -> Collection:
        """ 获取 MongoDB 集合实例 """

        return mongo_db[cls.collection_name]

    @classmethod
    def get_configured_indexes(cls) -> list[dict]:
        """ 获取配置的索引列表 """

        indexes: list[dict] = []

        # 单字段索引
        for field in cls.single_field_index:
            indexes.append(
                {
                    "name": f"{field}_1",
                    "key": [(field, ASCENDING)]
                }
            )

        # 联合字段索引
        for fields in cls.union_field_index:
            indexes.append(
                {
                    "name": f"{'_'.join(fields)}_1",
                    "key": [
                        (field, ASCENDING)
                        for field in fields
                    ]
                }
            )

        return indexes

    @classmethod
    def init_index(cls) -> None:
        """ 创建/更新/删除索引 """

        from components.logger import log

        collection = cls.get_collection()

        # 获取现有索引
        existing_indexes = {
            idx["name"]: idx
            for idx in collection.list_indexes()
            if idx["name"] != "_id_"  # 排除默认的_id索引
        }
        existing_names = set(existing_indexes.keys())

        # 获取配置的索引
        configured_indexes: list[dict] = cls.get_configured_indexes()
        configured_names: SetStr = {
            idx["name"]
            for idx in configured_indexes
        }

        # 创建或更新索引
        for idx in configured_indexes:
            idx_name = idx["name"]

            if idx_name not in existing_names:
                collection.create_index(
                    idx["key"],
                    name=idx_name,
                    background=True  # 后台创建索引，不阻塞读写操作
                )
                log.success(f"[MongoDB][索引] 创建 {collection.name}.{idx_name}")

            else:
                # 检查索引配置是否一致
                existing_idx = existing_indexes[idx_name]

                if existing_idx["key"] != idx["key"]:
                    collection.drop_index(idx_name)
                    collection.create_index(
                        idx["key"],
                        name=idx_name,
                        background=True
                    )
                    log.warning(f"[MongoDB][索引] 更新 {collection.name}.{idx_name}")

        # 删除不需要的索引
        for idx_name in existing_names - configured_names:
            collection.drop_index(idx_name)
            log.warning(f"[MongoDB][索引] 删除 {collection.name}.{idx_name}")

    @classmethod
    def create_index(cls) -> None:
        """ 创建索引 """

        sorting: int = ASCENDING  # 升序
        collection: Collection = cls.get_collection()

        for field in cls.single_field_index:
            collection.create_index({field: sorting})

        for fields in cls.union_field_index:
            collection.create_index(
                {
                    _field: sorting
                    for _field in fields
                }
            )

    @classmethod
    def get_deletion_query(cls, query_all: bool) -> dict[str, bool]:
        """ 数据是否删除的查询条件 """

        if query_all:
            return {}

        return cls.undeleted_query

    @classmethod
    def get_query(
        cls,
        query: DictOrNone,
        query_all: bool,
    ) -> dict:
        """ 获取查询条件 """

        return {
            **(query or {}),
            **cls.get_deletion_query(query_all),
        }

    @staticmethod
    def get_object_id_str(object_id: ObjectId | None) -> ObjectIdStr | None:
        """ 获取 ObjectId 字符串 """

        if object_id is None:
            return None

        return ObjectIdStr(object_id)

    @staticmethod
    def projection_shows(*args: str) -> DictStrInt | None:
        """
        需要展示的字段

        Args:
            args: 需要展示的字段

        Returns:
            需要展示的字段-字典
        """

        if not args:
            return None

        return {
            field: 1
            for field in args
        }

    @staticmethod
    def projection_no_shows(*args: str) -> DictStrInt | None:
        """
        不需要展示的字段

        Args:
            args: 需要展示的字段

        Returns:
            不需要展示的字段-字典
        """

        if not args:
            return None

        return {
            field: 0
            for field in args
        }

    @classmethod
    @inject_transaction_session
    def insert_one(
        cls,
        po: TPO,
        *args,
        **kwargs,
    ) -> InsertOneResult:
        """
        新增单个文档

        Args:
            po: PO 对象

        Returns:
            InsertOneResult
        """

        result: PyMongoInsertOneResult = cls.get_collection().insert_one(po.to_db_dict(), *args, **kwargs)

        return InsertOneResult(
            acknowledged=result.acknowledged,
            inserted_id=cls.get_object_id_str(result.inserted_id),
        )

    @classmethod
    @inject_transaction_session
    def insert_many(
        cls,
        po_list: Iterable[TPO],
        *args,
        **kwargs,
    ) -> InsertManyResult:
        """
        新增多个文档

        Args:
            po_list: 多个 PO 对象

        Returns:
            InsertManyResult
        """

        _data_list: Iterator[dict] = (
            po.to_db_dict()
            for po in po_list
        )

        result: PyMongoInsertManyResult = cls.get_collection().insert_many(_data_list, *args, **kwargs)

        return InsertManyResult(
            acknowledged=result.acknowledged,
            inserted_ids=[
                cls.get_object_id_str(_)
                for _ in result.inserted_ids
            ],
        )

    @classmethod
    @inject_transaction_session
    def insert_many_with_pagination(
        cls,
        po_list: Iterable[TPO],
        /,
        page_size: int = 500,
        *args,
        **kwargs,
    ) -> InsertManyResult:
        """
        分页新增多个文档

        Args:
            po_list: 多个PO 对象
            page_size: 分页参数，每次批量插入几个（默认为 500）

        Returns:
            InsertManyResult
        """

        acknowledged_list: list[bool] = []
        inserted_ids: list[ObjectIdStr] = []

        for chunk_pos in chunk_iterable(po_list, page_size):
            result: InsertManyResult = cls.insert_many(chunk_pos, *args, **kwargs)
            acknowledged_list.append(result.acknowledged)
            inserted_ids.extend(result.inserted_ids)

        return InsertManyResult(
            acknowledged=all(acknowledged_list),
            inserted_ids=inserted_ids,
        )

    @classmethod
    @inject_transaction_session
    def update_one(
        cls,
        po: TPO,
        now: DateTimeOrNone = None,
        *args,
        **kwargs,
    ) -> UpdateResult:
        """
        更新单个文档

        Args:
            po: PO 对象
            now: 当前时间（仅用于删除、批量删除、批量更新）

        Returns:
            UpdateResult
        """

        if kwargs.pop("delete", False):
            po.is_deleted = True

        po.update_time = now or DATETIME.now()
        _query: DictStrStr = {"id": po.id}
        _update: dict = {"$set": po.to_db_dict()}
        result: PyMongoUpdateResult = cls.get_collection().update_one(_query, _update, *args, **kwargs)

        return UpdateResult(
            acknowledged=result.acknowledged,
            matched_count=result.matched_count,
            modified_count=result.modified_count,
            upserted_id=cls.get_object_id_str(result.upserted_id),
        )

    @classmethod
    @inject_transaction_session
    def update_many(
        cls,
        po_list: Iterable[TPO],
        *args,
        **kwargs,
    ) -> UpdateResult:
        """
        更新多个文档

        Args:
            po_list: PO 对象列表

        Returns:
            UpdateResult
        """

        now: datetime = DATETIME.now()
        matched_count: int = 0
        modified_count: int = 0
        acknowledged_list: set[bool] = set()

        for po in po_list:
            result: UpdateResult = cls.update_one(po, now=now, *args, **kwargs)
            matched_count += result.matched_count
            modified_count += result.modified_count
            acknowledged_list.add(result.acknowledged)

        return UpdateResult(
            acknowledged=all(acknowledged_list),
            matched_count=matched_count,
            modified_count=modified_count,
        )

    @classmethod
    @inject_transaction_session
    def delete_one(
        cls,
        po: TPO,
        now: DateTimeOrNone = None,
        *args,
        **kwargs,
    ) -> UpdateResult:
        """
        删除单个文档

        Args:
            po: PO 对象
            now: 当前时间（仅用于批量删除）

        Returns:
            UpdateResult
        """

        return cls.update_one(po, now, delete=True, *args, **kwargs)

    @classmethod
    @inject_transaction_session
    def delete_many(
        cls,
        po_list: Iterable[TPO],
        *args,
        **kwargs,
    ) -> UpdateResult:
        """
        删除多个文档

        Args:
            po_list: PO 对象列表

        Returns:
            UpdateResult
        """

        return cls.update_many(po_list, delete=True, *args, **kwargs)

    @classmethod
    def find_one(
        cls,
        query: DictOrNone = None,
        projection: DictStrInt | None = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> TPO | None:
        """
        查询单个文档

        Args:
            query: 查询条件
            projection: 需要展示的字段
            query_all: 是否查询全部

        Returns:
            PO or None
        """

        _query: dict = cls.get_query(query, query_all)
        _projection: DictStrInt = projection or {}

        result: DocumentType | None = cls.get_collection().find_one(_query, projection=_projection, *args, **kwargs)

        if result is None:
            return None

        return cls.po(**result)

    @classmethod
    def find(
        cls,
        query: DictOrNone = None,
        projection: DictStrInt | None = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> Iterator[TPO]:
        """
        查询多个文档

        Args:
            query: 查询条件
            projection: 需要展示的字段
            query_all: 是否查询全部

        Returns:
            Iterator[PO]
        """

        _query: dict = cls.get_query(query, query_all)
        _projection: DictStrInt = projection or {}
        result: Cursor[DocumentType] = cls.get_collection().find(_query, projection=_projection, *args, **kwargs)

        for res in result:
            yield cls.po(**res)

    @classmethod
    def count_documents(
        cls,
        query: DictOrNone = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> int:
        """
        计算该集合中的文档数

        Args:
            query: 查询条件
            query_all: 是否查询全部

        Returns:
            查询到的文档数
        """

        _query: dict = cls.get_query(query, query_all)

        return cls.get_collection().count_documents(_query, *args, **kwargs)

    @classmethod
    def distinct(
        cls,
        field: str,
        query: DictOrNone = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> list:
        """
        获取去重后的字段值列表

        Args:
            field: 去重的字段名
            query: 查询条件
            query_all: 是否查询全部

        Returns:
            去重后的字段值列表
        """

        _query: dict = cls.get_query(query, query_all)

        return cls.get_collection().distinct(field, _query, *args, **kwargs)

    @classmethod
    def get_distinct_id_list(cls, id_list: Iterable[ObjectIdStr]) -> list[ObjectIdStr]:
        """
        获取去重后的 ID 列表

        Args:
            id_list: ID 列表

        Returns:
            去重后的 ID 列表
        """

        if not id_list:
            return []

        _query: dict = {
            "id": {
                "$in": list(id_list),
            },
        }

        id_list: list[str] = cls.distinct("id", _query)

        return [
            ObjectIdStr(_id)
            for _id in id_list
        ]

    @classmethod
    def paginated_query(
        cls,
        page_query: PageQuery,
        projection: DictStrInt | None = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> POPage[TPO]:
        """
        分页查询

        Args:
            page_query: 分页查询参数（已包括自定义查询条件）
            projection: 需要展示的字段
            query_all: 是否查询全部

        Returns:
            分页后的参数
        """

        _query: dict = cls.get_query(page_query.query, query_all)
        _projection: DictStrInt = projection or {}

        total_data_count: int = cls.count_documents(page_query.query, query_all, *args, **kwargs)
        """ 总数据条数 """

        total_page: int = math.ceil(total_data_count / page_query.page_size)
        """ 总页数 """

        skip_count: int = page_query.page_size * (page_query.page - 1)
        """ 当前查询需要跳过的数据条数 """

        raw_data: Cursor = cls.get_collection().find(_query, _projection, *args, **kwargs)
        """ 查询的原始数据：Cursor 对象，类似于迭代器，只能循环 1 次 """

        # 排序
        if page_query.sorting:
            raw_data = raw_data.sort(page_query.sorting)

        paged_data: Cursor = raw_data.skip(skip_count).limit(page_query.page_size)
        """ 分页后的数据 """

        po_list: Iterator[TPO] = (
            cls.po(**_data)
            for _data in paged_data
        )

        return POPage(
            page=page_query.page,
            page_size=page_query.page_size,
            page_count=total_page,
            data_count=total_data_count,
            data_list=po_list,
        )

    @classmethod
    def get_one_by_id(
        cls,
        id: ObjectIdStr | None,  # noqa: ignore[built-in name]
        projection: DictStrInt | None = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> TPO | None:
        """
        根据 ID 查询单个文档

        Args:
            id: ObjectId 字符串
            projection: 需要展示的字段
            query_all: 是否查询全部

        Returns:
            list[PO]
        """

        if not ObjectIdStr.is_valid(id):
            return None

        _query: dict = {"id": id}

        return cls.find_one(_query, projection, query_all, *args, **kwargs)

    @classmethod
    def exists(
        cls,
        query: dict,
        query_all: bool = False,
    ) -> bool:
        """
        根据条件查询文档是否存在

        Args:
            query: 查询条件
            query_all: 是否查询全部

        Returns:
            True / False
        """

        if not query:
            return False

        _query: dict = cls.get_query(query, query_all)
        _projection: DictStrInt = cls.projection_shows("_id")

        result: DocumentType | None = cls.get_collection().find_one(_query, projection=_projection)

        return result is not None

    @classmethod
    def exists_by_id(
        cls,
        id: ObjectIdStr | None,  # noqa: ignore[built-in name]
        query_all: bool = False,
    ) -> bool:
        """
        根据 ID 查询文档是否存在

        Args:
            id: ObjectId 字符串
            query_all: 是否查询全部

        Returns:
            True / False
        """

        if not ObjectIdStr.is_valid(id):
            return False

        return cls.exists({"id": id}, query_all=query_all)

    @classmethod
    def get_many_by_id_list(
        cls,
        id_list: Iterable[ObjectIdStr],
        projection: DictStrInt | None = None,
        query_all: bool = False,
        *args,
        **kwargs,
    ) -> Iterator[TPO]:
        """
        根据 ID 列表查询多个文档

        Args:
            id_list: ObjectId 列表
            projection: 需要展示的字段
            query_all: 是否查询全部

        Returns:
            返回一个迭代器，可能为空（如果 id_list 为空或查询无结果）
        """

        if not id_list:
            yield from ()  # 返回空生成器，类型是 Iterator[TPO]

        _query: dict = {
            "id": {
                "$in": list(id_list),
            },
        }

        yield from cls.find(_query, projection, query_all, *args, **kwargs)


TService = TypeVar("TService", bound=ServiceBase, covariant=True)
""" 继承 ServiceBase 的类（支持协变） """
